import ast
import platform
import warnings
from _ast import FunctionDef, AnnAssign, Name, Subscript, ClassDef, Attribute, Assign, Constant, AST
from ast import NodeVisitor

from typing import Any, Dict

import astunparse
import typing
from pprintast import pprintast

from MelodieAllegro.global_mgrs import id_generator
from MelodieAllegro.query.astquery import replace_all_names
from MelodieAllegro.types import PMType, string, integer, double, PMBasicType, PMArrayType, PMObjectType, \
    PMFunctionType, Modifier, PMStringType, PMIntegerSetType

from MelodieAllegro.utils import to_generic_id

PYTHON3_VERSION = int(platform.python_version_tuple()[1])


class UnsupportedSyntax(Exception):
    def __init__(self, node, err=""):
        self.text = f"Unsupported Syntax {astunparse.unparse(node)} at line {node.lineno}" + err

    def __str__(self):
        return self.text


class UndeclaredVariable(Exception):
    def __init__(self, node, var_name):
        self.text = f"Undefined variable {var_name} at line {node.lineno}. \n" \
                    f" Please define a variable by annotated assign like `i: int = 0`\n" + \
                    "Also, the variable should be defined at the beginning of a function, "

    def __str__(self):
        return self.text


class UnidentifiedAnnotation(Exception):
    def __init__(self, node, annotation_name: str):
        self.text = f"Cannot identify this Annotation `{annotation_name}` at line {node.lineno}."

    def __str__(self):
        return self.text


class UnTypedArgument(Exception):
    def __init__(self, node, method_name: str, arg_name: str):
        self.text = f"\n\nUntyped argument `{arg_name}` in function/method `{method_name}` at line {node.lineno}. \n" \
                    f"If this argument is not the `self` argument in object method, " \
                    f"please add annotation to this argument.\n" \
                    f"Just like: `def some_function(i)` -> `def some_function(i: int)`"

    def __str__(self):
        return self.text


def assert_isinstance(node: ast.AST, obj, typs, error: str = ""):
    if not isinstance(obj, typs):
        raise UnsupportedSyntax(node, error)


def is_valid_indexing(node: Subscript):
    if PYTHON3_VERSION <= 8:
        from ast import Index
        return isinstance(node.slice, Index) and isinstance(node.slice.value, (Name, Constant))
    else:
        return isinstance(node.slice, (Name, Constant))


def get_slice_value(node: Subscript):
    if PYTHON3_VERSION <= 8:
        from ast import Index
        return node.slice.value
    else:
        return node.slice


class TypeVisitor(NodeVisitor):
    def __init__(self, ):
        self.types: Dict[str, PMType] = {}  # 最好换成initial type 和 identified type
        self.identified_types: Dict[str, PMType] = {}

    def has_type(self, name: str):
        return (name in self.types) or (name in self.identified_types)

    def set_type(self, name, typ: PMType):
        self.identified_types[name] = typ

    def get_type(self, name) -> PMType:
        if name in self.types:
            return self.types[name]
        else:
            return self.identified_types[name]

    def visit_Name(self, node: Name) -> Any:
        default_types = {'str': string, 'int': integer, 'float': double}
        if node.id in ('int', 'float'):
            return PMBasicType(default_types[node.id])
        elif node.id == 'str':
            return PMStringType()
        elif node.id == 'IntSet':
            return PMIntegerSetType()
        elif self.has_type(node.id):
            return self.get_type(node.id)
        else:
            raise UnidentifiedAnnotation(node, annotation_name=node.id)

    def visit_Subscript(self, node: Subscript) -> Any:
        if isinstance(node.value, Name):
            # if isinstance(node.slice, Index) and isinstance(node.slice.value, Name):
            if is_valid_indexing(node):
                slice_value: typing.Union[Name, Constant] = get_slice_value(node)
                if node.value.id == 'Array':
                    val = self.visit(slice_value)
                    assert isinstance(val, PMType)
                    return PMArrayType(val)
                elif node.value.id in {"AgentList", 'List'}:
                    name = slice_value.id
                    if (name in {'float', 'int'}) or self.has_type(name):
                        return self.types[node.value.id + to_generic_id({'T': name})]
                    else:
                        raise NotImplementedError
                else:
                    raise NotImplementedError(node.value.id)
            else:
                raise UnsupportedSyntax(node)
        else:
            raise UnsupportedSyntax(node)

    def visit_AnnAssign(self, node: AnnAssign) -> Any:
        anno = node.annotation
        if isinstance(node.target, Name):
            if isinstance(anno, (Name, Subscript)):
                annotation_value = self.visit(anno)
                self.set_type(node.target.id, annotation_value)
            else:
                raise UnsupportedSyntax(node)
        else:
            raise UnsupportedSyntax(node)

    def start_visit(self, node):
        pprintast(node)
        return self.visit(node)


class ConstantTransformer(ast.NodeTransformer):
    """
    将数字类型的常量改为在entry块中定义，避免多次alloca造成段错误。

    将字符串类型的常量改成临时变量指针。

    这样，可以避免重复Malloc内存造成内存溢出。
    """

    def __init__(self, local_types=None):
        self.local_types: Dict[str, PMType] = {} if local_types is None else local_types
        self.const_definitions: typing.List[AnnAssign] = []

        # 已被定义的变量
        self.declared_vars = {}

    def start_visit(self, node):
        assert isinstance(node, FunctionDef)
        self.visit(node)
        for definition in self.const_definitions:
            node.body.insert(0, definition)

    def visit_AnnAssign(self, node: AnnAssign) -> Any:
        assert isinstance(node.target, Name)
        try:
            var_type = self.local_types[node.target.id]
        except KeyError:
            raise UndeclaredVariable(node, node.target.id)
        if isinstance(var_type, (PMBasicType, PMStringType)):
            self.const_definitions.append(
                AnnAssign(
                    target=node.target, value=var_type.get_ast_constant(), annotation=node.annotation, simple=1
                ))
            return Assign(targets=[node.target], value=node.value)
        else:
            return node

    def visit_Constant(self, node: Constant) -> Any:
        if isinstance(node.value, str):
            new_id = id_generator.new_id('const_str_literal')
            target = Name(id=new_id)
            self.const_definitions.append(AnnAssign(target=target, value=node, annotation=Name(id="str"), simple=1))
            self.local_types[new_id] = PMStringType()
            return target
        elif isinstance(node.value, (int, float)):
            return node
        else:
            raise UnsupportedSyntax(node)


class FuncTypeVisitor(NodeVisitor):
    def __init__(self, ):
        self._i = 1
        self.types: Dict[str, PMType] = {}

    def create_tmp_name(self):
        self._i += 1
        return '__tmp' + str(self._i)

    def visit_Name(self, node: Name) -> Any:
        return

    def visit_Assign(self, node: Assign) -> Any:
        assert len(node.targets) == 1
        target = node.targets[0]
        if isinstance(target, Name):
            assert target.id in self.types, f"line {node.lineno}: Variable {target.id} is not defined, " \
                                            f"please use annotated assign like `i: int = 0` to define a variable!"
        elif isinstance(target, Attribute):
            attr_target = self.visit(target.value)
            if isinstance(attr_target, Name):
                assert attr_target.id in self.types
            else:
                # 无需推导这类情况
                warnings.warn("Ignored attribute access with complexity")
                return
        elif isinstance(target, Subscript):
            attr_target = self.visit(target.value)
            if isinstance(attr_target, Name):
                assert attr_target.id in self.types
            else:
                warnings.warn("Ignored subscription with complexity")
                # 无需推导这类情况
                return
        else:
            raise NotImplementedError(target)

    def visit_AnnAssign(self, node: AnnAssign) -> Any:
        anno = node.annotation
        if isinstance(node.target, Name):
            if isinstance(anno, (Name, Subscript)):
                annotation_value = get_annotation_defs(anno, self.types)
                self.types[node.target.id] = annotation_value
            else:
                raise UnsupportedSyntax(node)
        else:
            raise UnsupportedSyntax(node)

    def visit(self, node):
        """Visit a node."""
        method = 'visit_' + node.__class__.__name__
        visitor = getattr(self, method, self.generic_visit)
        return visitor(node)

    def generic_visit(self, node):
        """Called if no explicit visitor function exists for a node."""
        for field, value in ast.iter_fields(node):
            if isinstance(value, list):
                for item in value:
                    if isinstance(item, AST):
                        self.visit(item)
            elif isinstance(value, AST):
                self.visit(value)

    def start_visit(self, node):
        pprintast(node)
        return self.visit(node)


def get_annotation_defs(anno_ast, types: Dict[str, PMType]):
    v = TypeVisitor()
    v.types = types.copy()
    return v.start_visit(anno_ast)


def scan_funcdef(func_ast, types: Dict[str, PMType], method_owner=None,
                 global_functions: Dict[str, PMFunctionType] = None,
                 ) -> PMFunctionType:
    assert isinstance(func_ast, FunctionDef)
    arg_types = []
    for arg in func_ast.args.args:
        if arg.annotation is not None:
            arg_type = get_annotation_defs(arg.annotation, types)
        else:
            try:
                arg_type = types[arg.arg]
            except KeyError:
                raise UnTypedArgument(func_ast, func_ast.name, arg.arg)
        assert isinstance(arg_type, PMType), arg_type
        types[arg.arg] = arg_type
        arg_types.append(arg_type)

    v = FuncTypeVisitor()
    v.types = types.copy()

    for item in func_ast.body:
        if isinstance(item, (AnnAssign, Assign)):
            getattr(v, 'visit_' + item.__class__.__name__)(item)
        else:
            warnings.warn(f"Ignored node {v} for no variable is defined/used !")

    if func_ast.returns is None:
        return_type = None
    else:
        return_visitor = TypeVisitor()
        return_visitor.types = types.copy()
        if isinstance(func_ast.returns, Name):

            rets = TypeVisitor.visit_Name(return_visitor, func_ast.returns)
        elif isinstance(func_ast.returns, Subscript):
            rets = TypeVisitor.visit_Subscript(return_visitor, func_ast.returns)
        else:
            raise NotImplementedError(func_ast.returns)
        return_type = rets

    local_types = v.types
    const_transformer = ConstantTransformer(local_types)
    const_transformer.start_visit(func_ast)

    f = PMFunctionType(func_ast.name, return_type,
                       arg_types, owner=method_owner,
                       is_method=method_owner is not None
                       )
    f.locals = local_types
    f.ast = func_ast

    print(f"Visited function decl {f.name} with args:", f.arg_types, "ret:", f.ret_type, "locals", f.locals)

    return f


def scan_class(filename: str, cls_name: str,
               all_classes: Dict[str, PMObjectType] = None,
               all_funcs: Dict[str, PMFunctionType] = None,
               generic_item: Dict[str, str] = None,
               ) -> PMObjectType:
    all_classes = {} if all_classes is None else all_classes
    all_funcs = {} if all_funcs is None else all_funcs
    if generic_item is None:
        generic_item = {"": ""}
    cls_ast = get_class_ast(filename, cls_name)
    if cls_ast is None:
        raise ValueError(f"Scan class error, maybe class `{cls_name}` not in `{filename}`")
    # if '' not in generic_item:
    #
    #     generic_name = '<' + ','.join([k + ' ' + v for k, v in generic_item.items()]) + ">"
    # else:
    #     generic_name = ''
    generic_name = to_generic_id(generic_item)
    print('generic_name', generic_name)
    cls_ast.name = cls_ast.name + generic_name
    cls_name = cls_ast.name
    replace_all_names(cls_ast, generic_item)

    ann_assigns = get_class_props(cls_ast)

    # process properties
    class_props: Dict[str, PMType] = {}
    if len(ann_assigns) == 0:
        warnings.warn(f"No property define phrase (like `prop_a: int`) found in class {cls_name}")

    parents: typing.List[PMObjectType] = []
    for parent_name in cls_ast.bases:
        if isinstance(parent_name, Name):

            assert parent_name.id in all_classes, f"Parent type `{parent_name.id}` is not defined!"
            parents.append(all_classes[parent_name.id])
        elif isinstance(parent_name, Subscript):
            assert isinstance(parent_name.value, Name)
            assert parent_name.value.id == 'Generic'
            pass
        else:
            raise NotImplementedError

    for ann_assign in ann_assigns:
        tv = TypeVisitor()
        tv.types.update(all_classes.copy())
        tv.start_visit(ann_assign)

        class_props.update(tv.identified_types)

    # process methods
    all_methods = get_all_methods(cls_ast)
    obj = PMObjectType(cls_name,
                       {k: (v, Modifier()) for k, v in
                        class_props.items()}
                       , parents=parents, methods=[])

    # 在解析funcdef参数时，必须用到的类型。
    funcdef_types = {'self': obj}
    funcdef_types.update(all_classes.copy())

    # build class constructor.
    # if '__init__' in all_methods:
    #     init_fcn = all_methods['__init__']
    #     init_fcn_def = scan_funcdef(init_fcn,
    #                                 funcdef_types,
    #                                 method_owner=obj
    #                                 )
    #     obj.methods[init_fcn_def.name] = init_fcn_def

    for method_name, method_ast in all_methods.items():
        a_fcn = query_function_ast(cls_ast, method_name)
        func_def = scan_funcdef(a_fcn,
                                funcdef_types,
                                method_owner=obj
                                )

        obj.child_methods[func_def.name] = func_def
        obj.methods[func_def.name] = func_def
    print(obj, obj.methods)
    return obj


def get_class_ast(filename, cls_name):
    with open(filename) as f:
        tree = ast.parse(f.read())
    for node in ast.walk(tree):
        if isinstance(node, ClassDef) and node.name == cls_name:
            return node


def query_function_ast(tree, func_name):
    for node in ast.walk(tree):
        if isinstance(node, FunctionDef) and node.name == func_name:
            return node
    raise ValueError(f"No function '{func_name}' found in ast.")


def get_node(tree, node_type):
    for node in ast.walk(tree):
        if isinstance(node, node_type):
            return node
    raise ValueError


def get_class_props(tree):
    cls_ast: ClassDef = get_node(tree, ClassDef)
    ann_assigns = []
    for item in cls_ast.body:
        if isinstance(item, AnnAssign):
            ann_assigns.append(item)
    return ann_assigns


def get_all_methods(tree):
    cls_ast: ClassDef = get_node(tree, ClassDef)
    method_defs = {}
    for item in cls_ast.body:
        if isinstance(item, FunctionDef):
            method_defs[item.name] = item
    return method_defs


def get_function_ast(filename, func_name):
    with open(filename) as f:
        tree = ast.parse(f.read())
    return query_function_ast(tree, func_name)
